#!/usr/bin/env python
import sys
import math
from sqlalchemy import create_engine, MetaData, Table, Column, String, ForeignKey, select, func
import csv
import settings as set
import itertools
import logging

def main(argv):
    #Set up logging
    logging.basicConfig(filename = set.logfile, level = logging.DEBUG)
    #Set up connection to population database
    xn, description, records = connectToDatabase()
    #Set up basic query
    q = setupQuery(xn, description, records) 
    #If requested, set up a additional dictionary to create an annotated
    #file
    if set.outfile:
        annotated = {}
        #open inputfile
        i = open(set.infile, 'r')
        reader = csv.DictReader(i, delimiter="\t")
        #save an entry for each patient
        for record in reader:
            if str(record) in annotated:
                annotated[str(record)]["count"] += 1
            else:
                annotated[str(record)] = record
                annotated[str(record)]["count"] = 1
        labels = reader.fieldnames
        i.close()

    expectedMatches = {}
    #For each experiment (vital information from experiment--fields for grouping)
    experiments = {"everything":Experiment(fields = [description.c.Race,
                           description.c.Sex,
                           description.c.minAge, description.c.maxAge,
                           records.c.state]),
                          "generalizedRace":Experiment(
                           race = lookupFn(set.races_generalized), 
                           fields = [description.c.Sex, description.c.minAge, 
                           description.c.maxAge, records.c.state]),
                          "fiveYears":Experiment(span = 5,
                          fields = [description.c.Race, description.c.Sex,
                          records.c.state]),
                          "tenYears":Experiment(span = 10,
                           fields = [description.c.Race, description.c.Sex,
                           records.c.state]),
                          "fiveGenRace":Experiment(span = 5,
                           race = lookupFn(set.races_generalized),
                           fields = [description.c.Sex,
                           records.c.state]),
                          "tenGenRace":Experiment(span = 10,
                           race = lookupFn(set.races_generalized),
                           fields = [description.c.Sex,
                           records.c.state]),
                          #"test":Experiment(
                          # span = 5,
                          # cutoff = 57,
                          # fields = [description.c.Race, description.c.Sex,
                          # records.c.state]),
                          "safeHarbor":Experiment(
                           cutoff = 89,
                           fields = [description.c.Race, description.c.Sex,
                           records.c.state])}
    for exptName, expt in experiments.items():
        logging.debug(exptName)
        #Do a preliminary scan of the records to group them
        QIDs = expt.scan(set.infile, xn, description, records, q) 
        #For each group
        expectedMatches[exptName] = []
        for g in QIDs:
            logging.debug(QIDs[g]["range"], g)
            for r in QIDs[g]["records"]:
                logging.debug("\t", r)  
        for group in QIDs:
            #Ensure that entire sample is covered
            if QIDs[group]["count"] < len(QIDs[group]["records"]):
                print "Had to adjust the Expected value of a bin from ",\
                      QIDs[group]["count"], " to ",\
                      len(QIDs[group]["records"]) 
                QIDs[group]["count"] = len(QIDs[group]["records"])
            #Set up a vector of size range of group
            bins = [0]*QIDs[group]["range"]
            #For each record in group
            for r in QIDs[group]["records"]:
                #Assign +1 to appropriate location in vector
                bins[r[0]] +=1
            #Add count - size of group / range of group
            expectedValue = (QIDs[group]["count"] - len(QIDs[group]["records"])) / \
                            (QIDs[group]["range"] * 1.)
            bins = [b + expectedValue for b in bins]    
            #Append risk to record
            expectedMatches[exptName].extend([bins[r[0]] for r in QIDs[group]["records"]])
            #OUTPUT PHASE
            logging.debug(QIDs[group]["range"], "bins", QIDs[group]["count"], "people", group)
            for r in QIDs[group]["records"]:
                if set.outfile:
                    annotated[str(r[1])][exptName] = bins[r[0]]
                logging.debug("\t", r, bins[r[0]])
        for r in expectedMatches[exptName]:
            assert r >= 1
    if set.outfile:
        w = open(set.outfile, 'w')
        labels.extend(experiments.keys()+["count"])
        out = csv.DictWriter(w, labels, delimiter="\t")
        out.writerow(dict(zip(labels, labels)))
        out.writerows(annotated.values())

    #Summary Phase
    w = open(set.sumfile, 'w')
    labels = ['exptName', 'totalPop', 'totalRisk']+\
             ['gDistinct'+str(t) for t in set.thresholds]+\
             ['trThresh'+str(t) for t in set.thresholds]
    dsWriter = csv.DictWriter(w, labels, delimiter="\t")
    dsWriter.writerow(dict(zip(labels, labels)))
    for expt, matches in expectedMatches.items():
        estimates={}
        estimates['exptName'] = expt
        estimates['totalPop'] = len(matches)
        estimates['totalRisk'] = sum([1./m for m in matches])
        for t in set.thresholds:
            estimates['gDistinct'+str(t)] = sum([1 for m in matches if m <= t])
            estimates['trThresh'+str(t)] = sum([1./m for m in matches if m <= t])      
        dsWriter.writerow(estimates) 
    w.close()

def lookupFn(dict):
    def lookup(word):
        if not word:
            return None
        return dict[word] if dict else word
    return lookup
def ageLookup(span, cutoff):
    def ageRange(age):
        if not age:
            return None 
        if cutoff and age > cutoff:
            minAge = cutoff + 1
            maxAge = -1
        else:
            minAge = age / span * span
            maxAge = (age / span + 1) * span - 1
            if cutoff and maxAge > cutoff:
                maxAge = cutoff
        return (minAge, maxAge, age)
    return ageRange
class Experiment():
    def __init__(self,
                 gender = lookupFn(set.genders),
                 race = lookupFn(set.races),
                 state = lookupFn(None),
                 span = 1,
                 cutoff = None,
                 fields = []):
        age = ageLookup(span, cutoff)
        self.transform={"gender":gender,
                        "race":race,
                        "state":state,
                        "age":age}
        self.span = span
        self.cutoff = cutoff
        self.fields = fields 
        self.QIDs = {}
    def scan(self, readfile, xn, description, records, query):
        #open readfile
        r = open(readfile, 'r')
        reader = csv.DictReader(r, delimiter="\t")
        #for line in file
        for record in reader:
            #get the demographic information
            demogs=getInfo(record)
            #transform info according to experiment
            demogs = dict([(k, self.transform[k](demogs[k])) for k in demogs])
            #print demogs
            #query population database
            result = generalQuery(xn, description, records, query, self.fields, demogs) 
            count = int(result["Pop"])
            if (not demogs["age"]) or \
                   (self.cutoff and result["groupMin"] > self.cutoff):
                ageSpan = 1
                position = 0
            else:
                ageSpan = (result["groupMax"] - result["groupMin"] + 1) / self.span
                if not ageSpan:
                    ageSpan = 1
                position = int(demogs["age"][2] - result["groupMin"])/self.span
            #add line to appropriate entry             
            if str(result) in self.QIDs:
                self.QIDs[str(result)]["records"].append((position, record))
            else:
                self.QIDs[str(result)] = {"range":ageSpan, "count":count, 
                                          "records": [(position, record)]}
            if "default" in result:
                self.QIDs[str(result)]["count"] = len(self.QIDs[str(result)]["records"])
        return self.QIDs            

def connectToDatabase():
    sqlengine = create_engine('%s://%s:%s@%s/%s'%\
                              (set.engine, set.uname, set.passwd,
                               set.hostAndPort, set.database))
    metadata = MetaData(sqlengine)
    metadata.create_all()
    xn = sqlengine.connect()

    description = Table('fieldNames_County', metadata, autoload = True)
    records = Table('allStates_County',
                    metadata,
                    Column('fieldname_id', 
                           String(10),
                           ForeignKey(description.c.FieldName)),
                    autoload=True)
    return xn, description, records

def getInfo(line):
    age = None
    birthYear = line[set.birthYearLabel]
    if birthYear:
        age = set.currentYear - int(birthYear)
    return {"gender":line[set.sexLabel],
            "race":line[set.raceLabel],
            "age":age}

def setupQuery(xn, description, records):
    s = select([description.c.minAge,
                description.c.maxAge,
                description.c.Race,
                description.c.Sex,
                func.min(description.c.minAge).label("groupMin"),
                func.max(description.c.maxAge).label("groupMax"),
                func.sum(records.c.population).label("Pop")],
               from_obj=[records.join(description)])
    return s

def generalQuery(xn, description, records, q, fields, demogs):
    q = q.where(records.c.state == set.state)
    q = q.where(description.c.Race.in_(demogs["race"])) if demogs["race"] else q
    q = q.where(description.c.Sex.in_(demogs["gender"])) if demogs["gender"] else q
    q = q.where(description.c.maxAge >= demogs["age"][0]) if demogs["age"] else q
    q = q.where(description.c.minAge <= demogs["age"][1]) \
            if (demogs["age"] and demogs["age"][1] > 0) \
            else q
    q = q.group_by(*fields)

    result = xn.execute(q)
    group = result.fetchall()
    try:
        match = group[0]
    except IndexError:
        print "Error for ", demogs,\
              "No matches found.  Returning a default match."
        match = {"groupMin":demogs["age"][0], "groupMax":demogs["age"][1],
                 "gender":demogs["gender"], "race":demogs["race"],
                 "Pop": 1, "default":True} 
    result.close()
    return match

if __name__ == '__main__':
    main(sys.argv[1:])
